package query

import (
	"fmt"
	"regexp"
	"strconv"
	"strings"
	"sync"
)

// EnflameQuery Enflame GPU查询实现
type EnflameQuery struct{}

func init() {
	RegisterQueryer("enflame", func() GPUQuery { return &EnflameQuery{} })
}

func (q *EnflameQuery) Query() ([]GPUInfo, error) {
	// 获取设备列表
	devices, err := q.getDevices()
	if err != nil {
		return nil, fmt.Errorf("获取设备列表失败: %v", err)
	}

	// 获取所有设备信息
	output, err := runCommand("efsmi", "-q", "-d", "MEMORY,USAGE")
	if err != nil {
		return nil, fmt.Errorf("获取设备信息失败: %v", err)
	}

	// 并行解析设备信息
	var wg sync.WaitGroup
	gpuInfos := make([]GPUInfo, len(devices))
	errors := make([]error, len(devices))

	for i, device := range devices {
		wg.Add(1)
		go func(index int, deviceId int) {
			defer wg.Done()
			utilization, memoryUsed, memoryTotal, err := q.parseDeviceInfo(output, deviceId)
			if err != nil {
				errors[index] = err
				return
			}
			gpuInfos[index] = GPUInfo{
				Index:       deviceId,
				Utilization: utilization,
				MemoryUsed:  memoryUsed,
				MemoryTotal: memoryTotal,
			}
		}(i, device)
	}
	wg.Wait()

	// 检查是否有错误
	for _, err := range errors {
		if err != nil {
			return nil, fmt.Errorf("解析设备信息失败: %v", err)
		}
	}

	return gpuInfos, nil
}

// getDevices 获取设备列表
func (q *EnflameQuery) getDevices() ([]int, error) {
	output, err := runCommand("efsmi", "-L")
	if err != nil {
		return nil, err
	}

	// 使用正则表达式快速提取设备ID
	deviceRegex := regexp.MustCompile(`^\s*(\d+)\s+`)
	lines := strings.Split(output, "\n")
	var devices []int

	for _, line := range lines {
		if matches := deviceRegex.FindStringSubmatch(line); matches != nil {
			if deviceID, err := strconv.Atoi(matches[1]); err == nil {
				devices = append(devices, deviceID)
			}
		}
	}

	if len(devices) == 0 {
		return nil, fmt.Errorf("未找到有效的GPU设备")
	}

	return devices, nil
}

// parseDeviceInfo 解析设备信息
func (q *EnflameQuery) parseDeviceInfo(output string, device int) (int, int, int, error) {
	// 预编译正则表达式
	deviceHeaderRegex := regexp.MustCompile(`DEV ID ` + strconv.Itoa(device))
	usageRegex := regexp.MustCompile(`GCU Usage\s*:\s*([\d.]+) %`)
	usedRegex := regexp.MustCompile(`Used Size\s*:\s*(\d+) MiB`)
	totalRegex := regexp.MustCompile(`Total Size\s*:\s*(\d+) MiB`)

	lines := strings.Split(output, "\n")
	inDeviceSection := false
	var utilization, used, total int
	foundUsage, foundUsed, foundTotal := false, false, false

	for _, line := range lines {
		// 检查是否进入设备部分
		if deviceHeaderRegex.MatchString(line) {
			inDeviceSection = true
			continue
		}

		// 如果不在设备部分，继续搜索
		if !inDeviceSection {
			continue
		}

		// 检查是否离开设备部分（遇到下一个设备）
		if strings.Contains(line, "DEV ID") && !strings.Contains(line, strconv.Itoa(device)) {
			break
		}

		// 解析使用率
		if !foundUsage {
			if matches := usageRegex.FindStringSubmatch(line); matches != nil {
				if util, err := strconv.ParseFloat(matches[1], 64); err == nil {
					utilization = int(util)
					foundUsage = true
				}
			}
		}

		// 解析已使用内存
		if !foundUsed {
			if matches := usedRegex.FindStringSubmatch(line); matches != nil {
				if value, err := strconv.Atoi(matches[1]); err == nil {
					used = value
					foundUsed = true
				}
			}
		}

		// 解析总内存
		if !foundTotal {
			if matches := totalRegex.FindStringSubmatch(line); matches != nil {
				if value, err := strconv.Atoi(matches[1]); err == nil {
					total = value
					foundTotal = true
				}
			}
		}

		// 如果所有信息都已找到，提前退出
		if foundUsage && foundUsed && foundTotal {
			break
		}
	}

	if !foundUsage || !foundUsed || !foundTotal {
		return 0, 0, 0, fmt.Errorf("未找到设备%d的完整信息", device)
	}

	return utilization, used, total, nil
}
